import numpy as np
import os
import output_function as of
import pandas as pd
import matplotlib.pyplot as plt

def vel_to_omega(u, v, dx, dy):  
    omega = (v[1:-1, 2:] - v[1:-1, :-2]) / 2 / dx - (u[2:, 1:-1] - u[:-2, 1:-1]) / 2 / dy  
    return omega

#扩散算符
def diff_opr1(myfun,nu,dx,dy):
    diff_y = nu * (myfun[2:, 1:-1] + myfun[:-2, 1:-1] - 2 * myfun[1:-1, 1:-1]) / dy / dy
    diff_x = nu * (myfun[1:-1, 2:] + myfun[1:-1, :-2] - 2 * myfun[1:-1, 1:-1]) / dx / dx
    return diff_x+diff_y



'''计算流函数，本质上是解泊松方程。要注意边界条件的处理'''
def compute_stream(omega_int, dx, dy, tol=1e-9, psi_init=None,max_iters=100000):
    """
    Jacobi iteration
    主输入:涡度场值omega_int
    辅输入:dx,dy,网格长,用于算差分矩阵
    主返回:计算得到的psi_init
    """
    if psi_init is None:
        psi_init = np.zeros((Ny, Nx))
    change = 1.
    iter = 0
    while (change > tol and iter<max_iters) or iter<10000:
        psi_new = psi_init.copy()
        #对中间区域格点的迭代过程，这里把差分矩阵在雅克比迭代情况的表达式写进来了。计算采用矩阵加法）
        psi_new[1:-1, 1:-1] = (omega_int \
                               + (psi_init[2:, 1:-1] + psi_init[:-2, 1:-1]) / dy / dy \
                               + (psi_init[1:-1, 2:] + psi_init[1:-1, :-2]) / dx / dx) / (2 / dx / dx + 2 / dy / dy)
        '''
        如果边界上流函数一直为零,初始时设成0就好,后面不动它
        顶盖第二类边条件(收敛性存在问题):psi_new[-1, 1:-1]=psi_new[-2, 1:-1]+dy
        '''
        #psi_new[-1, :]=psi_new[-2, :]   #顶部
        psi_new[-1, :]=0
        psi_new[0, :]=0  #底部
        psi_new[1:-1,0]=0
        psi_new[1:-1,-1]=0
        
        change = np.max(np.abs(psi_new - psi_init))  #change为迭代后矩阵元素变化最大值

        psi_init = psi_new
        iter += 1   
    end_signal = True if iter == 1 else False  #只用迭代1次时，end_signal返回真，这将用于停止迭代。因此此程序似乎在求从零直到稳态的过程。
    if iter>max_iters/10:
        print(iter)

    return psi_init, end_signal

# 使用示例
base_path = './data/RB_cavity_1e3/'  # 根据需要修改为您的保存路径
out_iter =5000  # 假设我们要加载第300时间步的数据

# 加载数据
n, psi, omega, u, v, out_iter_loaded = of.load_simulation_data(base_path, out_iter)
n_init0_matrix, xx, yy,x,y = of.load_initial_data(load_path=f'{base_path}initial_conditions.npz')

n_density=n

# 加载模拟参数
config = of.load_simulation_config(f'{base_path}simulation_config.json')
# 从config字典中提取特定的参数并赋值给变量
dx = config['dx']
dy = config['dy']
Lx = config['Lx']
Ly = config['Ly']
Nx = config['Nx']
Ny = config['Ny']
n_up = config['n_up']
n_0 = config['n_0']
n_down = config['n_down']
Delta_n = config['Delta_n']
Ra=Ra_star = config['Ra_star']
prandtl=Pr = config['Pr']
dt = config['dt']
ntime = config['ntime']
ndiag = config['ndiag']

# 打印一些变量以确认它们已被正确加载和赋值
print(f"dx: {dx}")
print(f"Ra_star: {Ra_star}")
print(f"dt: {dt}")
# 确认是否成功加载了正确的时间步长
if out_iter_loaded == out_iter:
    print(f"Successfully loaded data for time step: {out_iter_loaded}")
    print(omega.shape)
else:
    print("Failed to load the correct time step or file does not exist.")

# 接下来，您可以使用加载的数据进行分析、绘图或其他处理
t_total=out_iter*dt #总时间
#绘制流函数二维颜色图。
flag_contourf=True
if flag_contourf:
    fig, ax = plt.subplots()
    # 绘制等值线图，增加levels参数来使等值线更密集  
    levels = np.linspace(psi.min(), psi.max(), 10)  # 这里设置为50个级别  
    cs = ax.contourf(xx, yy, psi, levels=levels, cmap=plt.get_cmap('Spectral').reversed())  
    plt.streamplot(xx, yy, u, v, linewidth=0.5,color='black')
    # 添加标题  
    ax.set_title('Ra={},t={}'.format(Ra,t_total), fontsize=16)  # 设置标题文本和字体大小  
    ax.set_xlabel('X')  # 自行替换为合适的标签
    ax.set_ylabel('Y')  # 自行替换为合适的标签
    #添加colorbar
    cbar = fig.colorbar(cs)
    cbar.set_label('psi') 
    plt.show()

flag_contourf=True
if flag_contourf:
    fig, ax = plt.subplots()
    # 绘制等值线图，增加levels参数来使等值线更密集  
    levels = np.linspace(n.min(), n.max(), 20)  # 这里设置为50个级别  
    cs = ax.contourf(xx, yy, n, levels=levels, cmap=plt.get_cmap('Spectral').reversed())  
    # 添加标题  
    ax.set_title('Ra={},t={}'.format(Ra,t_total), fontsize=16)  # 设置标题文本和字体大小  
    # 添加x轴和y轴的标题
    ax.set_xlabel('X')  # 自行替换为合适的标签
    ax.set_ylabel('Y')  # 自行替换为合适的标签
    #添加colorbar
    cbar = fig.colorbar(cs)
    # 设置其标题
    cbar.set_label('T')  # 自行替换为合适的标签，例如'浓度'、'温度'等
    plt.show()

v_slice = v[Ny//2, :]
# 创建图表
plt.figure()
plt.plot(x, v_slice)
plt.xlabel('x')
plt.ylabel('v-midplane')  # 这里假设y轴是某种值，您可以根据实际情况修改
plt.title('Ra={},t={}'.format(Ra,t_total))
plt.show()